#include "ThreadPool.h"
#include <assert.h>
#include <memory>

namespace Ark
{
    struct ThreadPool::Task
    {
        Task(const std::function<void()> &cb)
            :callback(cb) {
        }
        Task(const std::function<void()> &&cb)
            :callback(cb) {
        }
        Task(const Task&) = delete;
        void operator = (const Task&) = delete;

        std::function<void()> callback;
    };

    ThreadPool::ThreadPool() {
    }

    ThreadPool::~ThreadPool() {
    }

    bool ThreadPool::Start(unsigned threadCount)
    {
        std::lock_guard<std::mutex> lock(_mutex);

        if (_running || !_threadList.empty()) {
            assert(!"already running");
            return false;
        }
        _running = true;
        _threadCount = threadCount != 0? threadCount : std::thread::hardware_concurrency();
        _StartWorkingThreads(_threadCount);
        return true;
    }

    void ThreadPool::_StartWorkingThreads(unsigned threadCount)
    {
        for (unsigned i = 0; i < threadCount; ++i) {
            std::shared_ptr<std::thread> thread(new std::thread(std::bind(&self::_WorkLoop, this)));
            _threadList.push_back(thread);
        }
    }

    bool ThreadPool::Stop()
    {
        {
            std::lock_guard<std::mutex> lock(_mutex);

            if (!_running || _threadList.empty()) {
                return true;
            }
            _running = false;
            _cond.notify_all(); // we should release lock after sending notification to work threads
        }

        for (auto t : _threadList) {
            t->join();
        }
        _threadList.clear();

        return true;
    }

    bool ThreadPool::Post(const std::function<void()>& cb) 
    {
        return _Post(new Task(cb));
    }

    bool ThreadPool::Post(const std::function<void()> &&cb)
    {
        return _Post(new Task(cb));
    }

    bool ThreadPool::_Post(ThreadPool::Task* task) 
    {
        std::lock_guard<std::mutex> lock(_mutex);

        if (!_running) {
            delete task;
            return false;
        }
        if (_taskList.empty()) {
            _cond.notify_one();
        }
        _taskList.push_back(task);
        return true;
    }

    size_t ThreadPool::GetTaskCount() const {
        std::lock_guard<std::mutex> lock(_mutex);
        return _taskList.size();
    }

    void ThreadPool::_WorkLoop()
    {
        while (_running)
        {
            std::unique_ptr<Task> work;

            {
                std::unique_lock<std::mutex> lock(_mutex);
                while (_taskList.empty()) {
                    if (!_running) {    // exit notification may be sent before lock is acquired
                        return;
                    }
                    _cond.wait(lock);
                }
                work.reset(_taskList.front());
                _taskList.pop_front();
            }
            
            work->callback();
        }
    }

}